import random,time
from typing import Optional
import numpy as np
import pandas as pd
from xml.dom.minidom import Element, parse

from .evdict import EVDict
from .ev import EV
from .trip import Trip

# 将0123转化为类型文字
def type2strplace(type): return ["Home","Work","Relax","Other"][type]

# 随机选一个不一样的
def random_diff(seq,exclude):
    ret=exclude
    while ret==exclude: ret=random.choice(seq)
    return ret

# 生成行程字典
def gendict_trip(*arr):
    return dict(zip(["trip_id","depart_time","from_TAZ","from_EDGE","to_TAZ","to_EDGE","routes","next_place_type"],arr))

# 从dict生成xml字符串
def genxml_trip(m,daynum=1):
    return ''.join([f'\n<trip id="{m["trip_id"]}" depart="{int(m["depart_time"])*60+86400*daynum}" ',
        f'fromTaz="{m["from_TAZ"]}" toTaz="{m["to_TAZ"]}" route_edges="{" ".join(m["routes"])}" />'])

class _EVinGen:
    def __init__(self,veh_id:str,vT,soc,v2g_prop:float):
        self.vehicle_id=veh_id
        self.bcap = vT["bcap_kWh"]
        self.soc = soc
        self.consump = vT["bcap_kWh"]/(vT["range_km"]*1000) #kWh/s
        self.efc_rate = vT["efc_rate_kW"]/3600.0 #kWh/s
        self.esc_rate = vT["esc_rate_kW"]/3600.0 #kWh/s
        self.max_v2g_rate = vT["max_V2G_kW"]/3600.0 #kWh/s
        self.omega = random.uniform(5.0,10.0)
        self.krel = random.uniform(1.0, 1.2)
        self.ksc = random.uniform(0.4, 0.6)
        self.kfc = random.uniform(0.2, 0.25)
        self.kv2g = random.uniform(0.65,0.75) if random.random()<v2g_prop else 1+1e-4
        self.trips=[]
        self.daynum=[]
    
    def add_trip(self,daynum:int,trip_dict:dict):
        self.daynum.append(daynum)
        self.trips.append(trip_dict)
    
    def to_xml(self)->str:
        ret=f'<vehicle id="{self.vehicle_id}" soc="{self.soc:.4f}" bcap="{self.bcap}" max_v2g_rate="{self.max_v2g_rate:.8f}" ' + \
            f'consump="{self.consump:.8f}" efc_rate="{self.efc_rate:.8f}" esc_rate="{self.esc_rate:.8f}" ' + \
            f'omega="{self.omega:.2f}" ksc="{self.ksc:.2f}" kfc="{self.kfc:.2f}" krel="{self.krel:.2f}" kv2g="{self.kv2g:.2f}">'
        for d,tr in zip(self.daynum,self.trips):
            ret+=genxml_trip(tr,d)
        ret+='\n</vehicle>\n'
        return ret
    
    def to_EV(self)->EV:
        trips=[Trip(m["trip_id"],int(m["depart_time"])*60+86400*daynum,m["routes"][0],m["routes"][-1]) 
               for m,daynum in zip(self.trips,self.daynum)]
        return EV(
            self.vehicle_id,
            trips,
            0.9,0.9,
            self.bcap,
            self.soc,
            self.consump,
            self.efc_rate,
            self.esc_rate,
            self.max_v2g_rate,
            self.omega,
            self.krel,
            self.kfc,
            self.ksc,
            self.kv2g,
        )
    
# 包装xml
class _xmlSaver:
    def __init__(self,path:str):
        self.a = open(path, "w")
        self.a.write("<root>\n")
    
    def write(self,e:_EVinGen):
        self.a.write(e.to_xml())
    
    def close(self):
        self.a.write("</root>")
        self.a.close()     
   
class EVsGenerator:
    def __init__(self,CROOT:str,PNAME:str,seed):
        '''
        CROOT: 行程参数文件夹
        PNAME: SUMO配置文件夹
        '''
        random.seed(seed)
        self.vTypes=pd.read_csv(CROOT+"\\车辆类型.csv").to_dict("records")
        #定义各功能区类型
        self.dic_taztype={}
        with open(PNAME+"/taz_type.txt","r") as fp:
            for ln in fp.readlines():
                name,lst=ln.split(":")
                self.dic_taztype[name.strip()]=[x.strip() for x in lst.split(',')]
        self.dic_taz={}
        root:Element = parse(PNAME+"/"+PNAME+".taz.xml").documentElement
        for taz in root.getElementsByTagName("taz"):
            if taz.hasAttribute("edges"):
                self.dic_taz[taz.getAttribute("id")]=taz.getAttribute("edges").split(' ')
            else:
                self.dic_taz[taz.getAttribute("id")]=[edge.getAttribute("id") for edge in taz.getElementsByTagName("tazSource")]

        #读取空间转移概率
        self.PSweekday = {}
        self.PSweekend = {}
        self.cdfweekday = {}
        self.cdfweekend = {}
        for dtype in ["Home","Work","Relax","Other"]:
            self.PSweekday[dtype]=pd.read_csv(f"{CROOT}\\4.空间转移概率\\{dtype}春天非周末.csv").drop(columns="Unnamed: 0")
            self.PSweekend[dtype]=pd.read_csv(f"{CROOT}\\4.空间转移概率\\{dtype}春天周末.csv").drop(columns="Unnamed: 0")
            self.cdfweekday[dtype]=pd.read_csv(f'{CROOT}\\2.停驻时长\\{dtype[0]}春天非周末.csv')
            self.cdfweekend[dtype]=pd.read_csv(f'{CROOT}\\2.停驻时长\\{dtype[0]}春天周末.csv')
        
        soc_pdf=pd.read_csv(f"{CROOT}\\0时SOC分布.csv")
        self.soc_vals=soc_pdf["range"].tolist()
        self.soc_freq=soc_pdf["freq"].tolist()

    def __getPs(self,is_weekday:bool,dtype:str)->pd.DataFrame:
        return self.PSweekday[dtype] if is_weekday else self.PSweekend[dtype]

    def __getcdf(self,is_weekday:bool,dtype:str)->pd.DataFrame:
        return self.cdfweekday[dtype] if is_weekday else self.cdfweekend[dtype]

    # 生成SOC
    def __genSoC(self):
        return random.choices(self.soc_vals,self.soc_freq)[0]/100.0
        
    def __getDest1(self,pfr:str,weekday:bool=True):
        '''
        获取第一次行程的下一次行程目的地
        :param pfr: 出发地功能区类型，例如“Home”
        :param weekday: 0/1变量。1表示工作日。0表示周末
        :return: 第一次行程返回：首次出行时间、到达目的地功能区类型，例如“Work”
        '''
        while True:
            while True:
                if weekday:
                    init_time = random.gammavariate(6.634511598855708, 65.76412399148424) + 114.54231929856397  # 非周末参数
                else:
                    init_time = random.gammavariate(3.45, 84.37) + 197.53 # 周末参数
                place = pfr
                init_time_i = int(init_time/15)
                ps=self.__getPs(weekday,place)
                if init_time_i<len(ps.columns): break
            data=ps.iloc[:,init_time_i]
            cdf = np.cumsum(data)
            if not cdf[3]==0: break
        x = random.random()
        next_place = type2strplace(cdf[cdf>x].index[0])
        return int(init_time),next_place

    def __getDestA(self,from_type,init_time_i,weekday:bool):
        '''
        获取*非*第一次行程的下一次行程目的地
        :param from_type:出发地类型，“Home”
        :param depart_time: 
        :return: 出发时间，目的地类型
        '''
        place = from_type
        data = self.__getPs(weekday,place).iloc[:,init_time_i]
        cdf = np.cumsum(data)*15
        if cdf[3]==0:
            to_type ="Home"
        else:
            x = random.random()
            to_type = type2strplace( cdf[cdf>x].index[0] )
        return to_type

    def __genFirstTrip1(self,trip_id,weekday:bool=True):
        '''
        生成首日的第1个行程
        :param trip_id:
        :return:trip的xml文件行、字典形式保存的相关信息
        dic_save ={"trip_id":trip_id,  行程id
                "depart_time":depart_time, #第一次行程出发时间/min
                "from_TAZ":from_TAZ,    #第一次行程出发区域TAZ类型“TAZ1”
                "from_EDGE":from_EDGE,   #出发道路边，例如“gnE29”
                "to_TAZ":to_TAZ,          #第一次行程到达区域TAZ类型“TAZ2”
                "to_EDGE":to_EDGE,        #到达道路边，例如“gnE2”
                "routes": routes,  #SUMO xml文件读取的routes类型'gneE22 gneE0 gneE16'
                "next_place_type":next_place_type}   #到达区域属性“Work”
        '''
        from_TAZ = random.choice(self.dic_taztype["Home"])  #出发TAZ区域编号，例如“TAZ1”
        from_EDGE = random.choice(self.dic_taz[from_TAZ]) #出发道路边，例如“gnE29”
        depart_time, next_place_type = self.__getDest1("Home", weekday) #得到出发时间和目的地区域类型
        to_TAZ = random.choice(self.dic_taztype[next_place_type])#到达TAZ区域编号，例如“TAZ2”
        to_EDGE = random_diff(self.dic_taz[to_TAZ],from_EDGE)  #到达TAZ边，例如“gnE28”
        return gendict_trip(trip_id,depart_time,from_TAZ,from_EDGE,to_TAZ,to_EDGE,[from_EDGE,to_EDGE],next_place_type)

    cdf_dict={}

    def __genStopTime(self,from_type,weekday:bool):
        cdf = np.cumsum(self.__getcdf(weekday,from_type)["0"])*15
        return cdf[cdf > random.random()].index[0] + 1  # 停留时长单位为15min   #抽取停留时间

    def __genTripA(self,trip_id,from_TAZ,from_type,from_EDGE,start_time,weekday:bool=True):
        '''
        生成第2个行程
        '''
        stop_duration = self.__genStopTime(from_type,weekday)
        depart_time = start_time+stop_duration*15+20
        next_place2 = self.__getDestA(from_type, stop_duration, weekday)
        taz_choose2 = random.choice(self.dic_taztype[next_place2])
        edge_choose2 = random_diff(self.dic_taz[taz_choose2],from_EDGE)
        return gendict_trip(trip_id,depart_time,from_TAZ,from_EDGE,taz_choose2,edge_choose2,[from_EDGE,edge_choose2],next_place2)

    def __genTripF(self,trip_id,from_TAZ,from_type,from_EDGE,start_time,first_TAZ,first_EDGE,weekday:bool=True):
        '''
        生成第3个行程
        '''
        stop_time = self.__genStopTime(from_type,weekday)
        depart_time = start_time+stop_time+20
        return gendict_trip(trip_id,depart_time,from_TAZ,from_EDGE,first_TAZ,first_EDGE,[from_EDGE,first_EDGE],"Home")
            
    def __genTripsChain1(self,vehicle_id:str,v2g_prop:float):#vehicle_trip
        '''
        生成首日一整天的出行链
        :param vehicle_id: 车辆ID
        :param v2g_prop: 愿意V2G的概率
        :return:
        '''
        daynum=0
        weekday=True
        ev=_EVinGen(vehicle_id,random.choice(self.vTypes),self.__genSoC(),v2g_prop)
        trip_1 = self.__genFirstTrip1("trip0_1", weekday)
        ev.add_trip(daynum,trip_1)
        trip_2 = self.__genTripA("trip0_2",trip_1["to_TAZ"],trip_1["next_place_type"],trip_1["to_EDGE"],trip_1["depart_time"],weekday)
        ev.add_trip(daynum,trip_2)
        trip_3 = self.__genTripF("trip0_3", trip_2["to_TAZ"], trip_2["next_place_type"],trip_2["to_EDGE"],
            trip_2["depart_time"], trip_1["from_TAZ"],trip_1["routes"][0], weekday)
        if trip_3["routes"][0]!=trip_3["routes"][-1]: ev.add_trip(daynum,trip_3)
        return ev

    def __genFirstTripA(self,trip_id,ev:_EVinGen,weekday:bool=True):
        '''
        生成非首日的第1段行程
        :param trip_id: 行程id
        :param vehicle_node: 车辆节点如rootNode.getElementsByTagName("vehicle")[0]
        :param weekday: 是否是工作日
        :return:
        '''
        trip_last = ev.trips[-1]
        from_EDGE = trip_last["routes"][-1]
        from_TAZ = trip_last["to_TAZ"]
        [depart_time, next_place_type] = self.__getDest1("Home", weekday)  # 得到出发时间和目的地区域类型
        to_TAZ = random.choice(self.dic_taztype[next_place_type])  # 到达TAZ区域编号，例如“TAZ2”
        to_EDGE = random_diff(self.dic_taz[to_TAZ],from_EDGE) # 到达TAZ边，例如“gnE28”
        return gendict_trip(trip_id,depart_time,from_TAZ,from_EDGE,to_TAZ,to_EDGE,[from_EDGE,to_EDGE],next_place_type)

    def __genTripsChainA(self,ev:_EVinGen, daynum:int=1):#vehicle_trip
        '''
        生成非首日一整天的出行链
        '''
        weekday = (daynum-1)%7+1 in [1,2,3,4,5]
        trip2_1 = self.__genFirstTripA(f"trip{daynum}_1",ev,weekday)
        ev.add_trip(daynum,trip2_1)
        trip2_2 = self.__genTripA(f"trip{daynum}_2",trip2_1["to_TAZ"],trip2_1["next_place_type"],trip2_1["to_EDGE"],trip2_1["depart_time"],weekday)
        ev.add_trip(daynum,trip2_2)
        trip2_3 = self.__genTripF(f"trip{daynum}_3", trip2_2["to_TAZ"], trip2_2["next_place_type"],trip2_2["to_EDGE"],
            trip2_2["depart_time"], trip2_1["from_TAZ"],trip2_1["routes"][0],weekday)
        if trip2_3["routes"][0]!=trip2_3["routes"][-1]: ev.add_trip(daynum,trip2_3)

    def __genEV(self,veh_id:str,v2g_prop:float)->_EVinGen:
        ev = self.__genTripsChain1(veh_id,v2g_prop)
        for j in range(1,8):
            self.__genTripsChainA(ev,j)
        return ev
    
    def genEV(self,veh_id:str,v2g_prop:float)->EV:
        '''
        生成一辆车一整周的出行链
        '''
        return self.__genEV(veh_id,v2g_prop).to_EV()

    def genEVs(self,N:int,v2g_prop:float,fname:Optional[str]=None,silent:bool=False)->EVDict:
        '''
        生成N辆车并以指定格式保存至文件fname(为None则不保存), 同时显示进度, 并以EVDict的形式返回
        '''
        st_time=time.time()
        saver=_xmlSaver(fname) if fname is not None else None
        last_time=0
        ret=EVDict()
        for i in range(0,N):
            ev = self.__genEV("v"+str(i),v2g_prop)
            ret.add(ev.to_EV())
            if saver is not None: saver.write(ev)
            this_time=time.time()
            if this_time>=last_time+1:
                if not silent: print(f"\r{i}/{N}, {i/N*100:.2f}%  ",end=" ")
                last_time=this_time
        if not silent: print(f"\r已完成. 用时:{time.time()-st_time:.1f}秒     ")
        if saver is not None: saver.close()
        return ret
